import re
import collections
import numpy as np
import tensorflow as tf
import modeling
import tokenization


class Encode(object):
    def __init__(self, bert_config, vocab_file):
        self.bert_config = bert_config
        self.tokenizer = tokenization.FullTokenizer(vocab_file=vocab_file, do_lower_case=False)

        self._add_ops()
        self._output()

    def _add_ops(self):
        self.input_ids = tf.placeholder(tf.int32, shape=[None, None], name='input_ids')
        self.input_mask = tf.cast(self.input_ids > 0, tf.int32, name='input_mask')

    def _output(self):
        model = modeling.BertModel(config=self.bert_config,
                                   is_training=False,
                                   input_ids=self.input_ids,
                                   input_mask=self.input_mask
                                    )
        self.output = model.get_sequence_output()

    def init(self, checkpoint_path):
        tvars = tf.trainable_variables()
        assignment_map, initialized_variable_names = get_assignment_map_from_checkpoint(tvars, checkpoint_path)
        tf.train.init_from_checkpoint(checkpoint_path, assignment_map)

    def get_feature(self, sess, inputs):
        feed_dict = {self.input_ids: inputs}

        return sess.run(self.output, feed_dict)

    def encode(self, sess, query):
        tokens = self.tokenizer.tokenize(query)
        tokens = ['[CLS]'] + tokens + ['[SEP]']
        input_ids = self.tokenizer.convert_tokens_to_ids(tokens)
        input_ids = np.reshape(input_ids, (1, -1))

        return self.get_feature(sess, input_ids)


def get_assignment_map_from_checkpoint(tvars, init_checkpoint):
  """Compute the union of the current variables and checkpoint variables."""
  assignment_map = {}
  initialized_variable_names = {}

  name_to_variable = collections.OrderedDict()
  for var in tvars:
    name = var.name
    m = re.match("^(.*):\\d+$", name)
    if m is not None:
      name = m.group(1)
    name_to_variable[name] = var

  init_vars = tf.train.list_variables(init_checkpoint)

  assignment_map = collections.OrderedDict()
  for x in init_vars:
    (name, var) = (x[0], x[1])
    if name not in name_to_variable:
      continue
    assignment_map[name] = name
    initialized_variable_names[name] = 1
    initialized_variable_names[name + ":0"] = 1

  return (assignment_map, initialized_variable_names)
